import tkinter as tk
from tkinter import ttk, filedialog, scrolledtext
import paramiko
import threading
import time
import re
import logging
from openpyxl import load_workbook

logger = logging.getLogger(__name__)


class ServerConnection:
    def __init__(self, main_window):
        self.main_window = main_window
        self.ssh_sessions = {}  # 活动会话
        self.ssh_cache = {}  # 所有缓存的会话
        self.terminal_outputs = {}  # 终端输出文本框字典

    def create_server_table(self):
        """创建服务器表格"""
        table_frame = ttk.Frame(self.main_window.top_frame, padding="10 5 10 3")
        table_frame.pack(fill=tk.BOTH, expand=True)

        # 创建表格
        columns = ("选择", "服务器IP", "端口", "用户名", "密码", "连接状态", "服务器架构", "命令执行状态",
                   "文件上传状态","工控机主机名","主工控机IP","备工控机IP","工控机虚拟IP","去重服务器IP")
        style = ttk.Style()
        style.map("Treeview",
                  font=[('selected', ('Arial', 11))])
        self.server_table = ttk.Treeview(table_frame, columns=columns, show="headings",height=5)

        # 设置列宽和标题
        for col in columns:
            self.server_table.heading(col, text=col)
            width = 100
            if col == "选择":
                width = 40
            elif col in ("端口", "连接状态", "服务器架构"):
                width = 80
            elif col in ("命令执行状态", "文件上传状态"):
                width = 150
            self.server_table.column(col, width=width)

        # 添加滚动条
        scrollbar = ttk.Scrollbar(table_frame, orient=tk.VERTICAL, command=self.server_table.yview)
        scrollbar_x = ttk.Scrollbar(table_frame, orient=tk.HORIZONTAL, command=self.server_table.xview)
        self.server_table.configure(yscroll=scrollbar.set, xscroll=scrollbar_x.set)
        scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        scrollbar_x.pack(side=tk.BOTTOM, fill=tk.X)
        self.server_table.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        # 绑定选择事件
        self.server_table.bind("<ButtonRelease-1>", self.toggle_selection)

    def toggle_selection(self, event):
        """切换选择状态"""
        region = self.server_table.identify_region(event.x, event.y)
        if region == "cell":
            column = self.server_table.identify_column(event.x)
            if column == "#1":  # 选择列
                item = self.server_table.identify_row(event.y)
                if item:
                    values = list(self.server_table.item(item, "values"))
                    values[0] = "√" if values[0] == "" else ""
                    self.server_table.item(item, values=values)
                    self._update_select_all_button_text()

    def _update_select_all_button_text(self):
        """根据表格选择状态更新全选/取消全选按钮的文本"""
        all_selected = True
        has_any_selection = False
        for item in self.server_table.get_children():
            values = self.server_table.item(item, "values")
            if values[0] != "√":
                all_selected = False
            else:
                has_any_selection = True

        if all_selected and has_any_selection:
            self.main_window.select_all_button_text.set("取消全选")
        else:
            self.main_window.select_all_button_text.set("全选")

    def toggle_select_all(self):
        """全选/取消全选"""
        all_selected = True

        # 检查是否所有行都已选中
        for item in self.server_table.get_children():
            values = self.server_table.item(item, "values")
            if values[0] != "√":
                all_selected = False
                break

        # 根据状态选择或取消全选所有行
        for item in self.server_table.get_children():
            values = list(self.server_table.item(item, "values"))
            values[0] = "" if all_selected else "√"
            self.server_table.item(item, values=values)

        # 更新按钮文本
        self._update_select_all_button_text()

    def load_server_info(self):
        """加载服务器信息"""
        file_path = self.main_window.server_file_path.get()
        if not file_path:
            self.main_window.show_message("请选择服务器信息文件")
            return

        try:
            # 清空表格
            for item in self.server_table.get_children():
                self.server_table.delete(item)

            if file_path.endswith('.xlsx'):
                # 用 openpyxl 读取 Excel 文件
                wb = load_workbook(file_path)
                ws = wb.active
                headers = [cell.value for cell in next(ws.iter_rows(min_row=1, max_row=1))]
                data_rows = ws.iter_rows(min_row=2, values_only=True)
                total_rows = ws.max_row - 1
            elif file_path.endswith('.txt'):
                # 读取 TXT 文件
                with open(file_path, 'r', encoding='utf-8') as f:
                    lines = f.readlines()
                if not lines:
                    self.main_window.show_message("TXT文件为空")
                    return
                
                # 首行为列名，去除头尾空格
                headers = [h.strip() for h in lines[0].strip().split(',')]
                data_rows = []
                for line in lines[1:]:
                    if not line.strip(): # 跳过空行
                        continue
                    # 数据从第2行开始，列之间用逗号分隔，去除字段头尾空格
                    data_rows.append([d.strip() for d in line.strip().split(',')])
                total_rows = len(data_rows)
            else:
                self.main_window.show_message("不支持的文件格式，请选择.xlsx或.txt文件")
                return

            for row in data_rows:
                # 跳过所有字段都为空的行
                if not any(row):
                    continue
                # 确保行数据与头部长度匹配，不匹配则跳过或报错
                if len(row) != len(headers):
                    self.main_window.show_message(f"警告: 行数据与列头不匹配，跳过此行: {row}")
                    continue

                row_dict = dict(zip(headers, row))
                # 使用.get()方法安全获取字典值，提供默认空字符串以防键不存在
                values = (
                    "",
                    row_dict.get('服务器IP', ''),
                    row_dict.get('端口', ''),
                    row_dict.get('用户名', ''),
                    row_dict.get('密码', ''),
                    "未连接",
                    "", "", "", # 额外字段，如果TXT中没有，则为空
                    row_dict.get("工控机主机名", ''),
                    row_dict.get("主工控机IP", ''),
                    row_dict.get("备工控机IP", ''),
                    row_dict.get("工控机虚拟IP", ''),
                    row_dict.get("去重服务器IP", '')
                )
                item_id = self.server_table.insert("", tk.END, values=values)

            self.main_window.show_message(f"成功加载 {total_rows} 条服务器信息")
            self._update_select_all_button_text()

        except Exception as e:
            self.main_window.show_message(f"加载服务器信息失败: {str(e)}")

    def connect_selected_servers(self):
        """连接选中的服务器"""
        logger.info("尝试连接选中的服务器...")
        selected_servers = []
        for item in self.server_table.get_children():
            values = self.server_table.item(item, "values")
            if values[0] == "√":
                server_info = {
                    'id': item,
                    'ip': values[1],
                    'port': int(values[2]),
                    'username': values[3],
                    'password': values[4]
                }
                selected_servers.append(server_info)

        if not selected_servers:
            self.main_window.show_message("请选择要连接的服务器")
            logger.warning("连接服务器：没有选中的已连接服务器")
            return

        # 多线程连接服务器
        for server in selected_servers:
            threading.Thread(target=self.connect_server, args=(server,)).start()

        self._update_select_all_button_text()

    def connect_server(self, server_info):
        """连接单个服务器"""
        item_id = server_info['id']
        ip = server_info['ip']
        port = server_info['port']
        username = server_info['username']
        password = server_info['password']

        logger.info(f"尝试连接到 {ip}:{port} (用户: {username})")

        # 检查是否已有缓存的连接
        if ip in self.ssh_cache and self.ssh_cache[ip]['ssh'] is not None:
            try:
                # 尝试发送空命令来测试连接是否仍然有效
                shell = self.ssh_cache[ip]['shell']
                shell.send("\n")
                time.sleep(0.1)
                shell.recv(1024)  # 清除输出缓冲区

                # 连接仍然有效，使用缓存
                logger.info(f"使用缓存重新连接到 {ip}")
                self.main_window.show_message(f"使用缓存重新连接到 {ip}")

                # 更新活动会话列表
                self.ssh_sessions[ip] = self.ssh_cache[ip]

                # 更新连接状态
                self.update_server_status(item_id, "已连接(缓存)", 5)

                # 确保该服务器有对应的终端标签页
                self.ensure_server_tab(ip)

                return
            except Exception as e:
                logger.warning(f"缓存的连接 {ip} 已失效: {e}, 尝试重新建立连接")
                self.main_window.show_message(f"缓存的连接已失效，重新连接到 {ip}")

        try:
            # 更新连接状态
            self.update_server_status(item_id, "连接中...", 5)
            logger.info(f"正在建立新的SSH连接到 {ip}...")

            # 创建SSH客户端
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

            jump_ssh = None

            if self.main_window.enable_ssh_tunnel.get():
                jump_host = self.main_window.jump_host_server.get()
                jump_username = self.main_window.jump_host_username.get()
                jump_password = self.main_window.jump_host_password.get()

                if not jump_host or not jump_username or not jump_password:
                    self.update_server_status(item_id, "隧道信息不完整", 5)
                    self.main_window.show_message(f"连接 {ip} 失败: SSH隧道信息不完整")
                    logger.error(f"连接 {ip} 失败: SSH隧道信息不完整")
                    return

                try:
                    self.main_window.show_message(f"正在通过踏板机 {jump_host} 建立SSH隧道...")
                    jump_ssh = paramiko.SSHClient()
                    jump_ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                    jump_ssh.connect(jump_host, username=jump_username, password=jump_password, timeout=10)

                    # 创建一个socket，通过跳板机连接目标服务器
                    transport = jump_ssh.get_transport()
                    dest_addr = (ip, port)
                    local_addr = ('', 0)
                    channel = transport.open_channel("direct-tcpip", dest_addr, local_addr)

                    # 使用这个channel作为SSHClient的socket
                    ssh.connect(ip, port, username, password, sock=channel, timeout=10)
                    logger.info(f"成功通过踏板机 {jump_host} 连接到 {ip}")
                    self.main_window.show_message(f"成功通过踏板机 {jump_host} 连接到 {ip}")
                except Exception as e:
                    self.update_server_status(item_id, "隧道连接失败", 5)
                    self.main_window.show_message(f"通过踏板机连接 {ip} 失败: {str(e)}")
                    logger.error(f"通过踏板机连接 {ip} 失败: {str(e)}", exc_info=True)
                    return
            else:
                # 不使用隧道，直接连接
                ssh.connect(ip, port, username, password, timeout=10)

            # 创建交互式Shell
            shell = ssh.invoke_shell()
            shell.settimeout(10)

            # 等待初始提示符
            time.sleep(1)
            init_output = shell.recv(4096).decode("utf-8", errors="ignore")

            # 创建或激活对应的标签页
            self.ensure_server_tab(ip)
            self.append_to_terminal(ip, f"已连接到 {ip}\n{init_output}")

            # 在shell中执行命令获取服务器架构
            shell.send("uname -m\n")
            time.sleep(1)
            output = shell.recv(4096).decode("utf-8", errors="ignore")

            # 提取架构信息
            arch_match = re.search(r'(x86_64|aarch64|armv\w+)', output)
            architecture = arch_match.group(1) if arch_match else "未知"

            # 检查是否需要develop模式
            if self.main_window.develop_mode.get():
                dev_password = self.main_window.develop_password.get()
                if not dev_password:
                    self.update_server_status(item_id, "缺少develop密码", 5)
                    self.main_window.show_message(f"连接 {ip} 失败: 缺少develop密码")
                    logger.error(f"连接 {ip} 失败: 缺少develop密码")
                    ssh.close()
                    return

                # 发送develop命令
                self.append_to_terminal(ip, "\n$ develop\n")
                shell.send("develop\n")
                time.sleep(1)
                output = shell.recv(4096).decode("utf-8", errors="ignore")
                self.append_to_terminal(ip, output)

                # 如果有密码提示，发送密码
                if "Password:" in output or "password:" in output.lower():
                    logger.info(f"发送develop密码到 {ip}")
                    shell.send(f"{dev_password}\n")
                    time.sleep(1)
                    output = shell.recv(4096).decode("utf-8", errors="ignore")
                    self.append_to_terminal(ip, output)

                    if "incorrect" in output.lower() or "denied" in output.lower():
                        self.update_server_status(item_id, "develop密码错误", 5)
                        logger.error(f"连接 {ip} 失败: develop密码错误")
                    else:
                        self.update_server_status(item_id, "已连接(develop)", 5)
                        logger.info(f"成功连接到 {ip} (develop模式)")
                else:
                    self.update_server_status(item_id, "develop模式失败", 5)
                    logger.error(f"连接 {ip} 失败: 未收到密码提示，develop模式可能失败")
            else:
                self.update_server_status(item_id, "已连接", 5)
                logger.info(f"成功连接到 {ip}")

            # 存储SSH会话到活动会话和缓存
            session_data = {
                'ssh': ssh,
                'shell': shell,
                'item_id': item_id,
                'jump_ssh': jump_ssh
            }
            self.ssh_sessions[ip] = session_data
            self.ssh_cache[ip] = session_data

            # 更新服务器架构
            self.update_server_architecture(item_id, architecture)

            # 切换到该服务器的标签页
            self.switch_to_server_tab(ip)

        except Exception as e:
            error_msg = str(e)
            self.update_server_status(item_id, f"连接失败: {error_msg[:20]}...", 5)
            self.main_window.show_message(f"连接 {ip} 失败: {error_msg}")
            logger.error(f"连接 {ip} 失败: {error_msg}", exc_info=True)

    def ensure_server_tab(self, ip):
        """确保服务器有对应的标签页，如果没有则创建"""
        tab_exists = False

        # 检查标签页是否已存在
        for i in range(self.main_window.terminal_execution.terminal_notebook.index("end")):
            if self.main_window.terminal_execution.terminal_notebook.tab(i, "text") == f"终端-{ip}":
                tab_exists = True
                break

        if not tab_exists:
            # 创建新标签页
            server_tab = ttk.Frame(self.main_window.terminal_execution.terminal_notebook)
            self.main_window.terminal_execution.terminal_notebook.add(server_tab, text=f"终端-{ip}")

            # 创建文本输出框
            output_text = scrolledtext.ScrolledText(server_tab, bg="black", fg="white", wrap=tk.WORD, height=8)
            output_text.pack(fill=tk.BOTH, expand=True)
            output_text.config(state=tk.DISABLED)

            # 存储文本框引用
            self.terminal_outputs[ip] = output_text

    def switch_to_server_tab(self, ip):
        """切换到指定服务器的标签页"""
        for i in range(self.main_window.terminal_execution.terminal_notebook.index("end")):
            if self.main_window.terminal_execution.terminal_notebook.tab(i, "text") == f"终端-{ip}":
                self.main_window.terminal_execution.terminal_notebook.select(i)
                break

    def append_to_terminal(self, ip, text):
        """向指定服务器的终端标签页添加文本"""
        # ANSI转义序列正则表达式
        ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])')
        text = ansi_escape.sub('', text)
        if ip in self.terminal_outputs:
            output_text = self.terminal_outputs[ip]
            output_text.config(state=tk.NORMAL)
            output_text.insert(tk.END, text)
            output_text.see(tk.END)
            output_text.config(state=tk.DISABLED)

    def update_server_status(self, item_id, status, column):
        """更新服务器在表格中的状态"""
        values = list(self.server_table.item(item_id, "values"))
        values[column] = status
        self.server_table.item(item_id, values=values)

    def update_server_architecture(self, item_id, architecture):
        """更新服务器架构信息"""
        values = list(self.server_table.item(item_id, "values"))
        values[6] = architecture
        self.server_table.item(item_id, values=values)

    def disconnect_selected_servers(self):
        """断开选中的服务器连接"""
        logger.info("尝试断开选中的服务器连接...")
        # 获取选中的服务器
        for item in self.server_table.get_children():
            values = self.server_table.item(item, "values")
            if values[0] == "√" and values[5] in ["已连接", "已连接(develop)", "已连接(缓存)"]:
                ip = values[1]
                if ip in self.ssh_sessions:
                    try:
                        session = self.ssh_sessions[ip]
                        
                        # 1. Close the main SSH connection (which uses the tunnel channel)
                        if session['ssh'] is not None:
                            session['ssh'].close()
                            logger.info(f"已关闭 {ip} 的SSH连接")

                        # 2. Then, close the jump host connection if it exists
                        if 'jump_ssh' in session and session['jump_ssh']:
                            session['jump_ssh'].close()
                            logger.info(f"已关闭 {ip} 的跳板机连接")

                        # 从活动会话和缓存中移除
                        del self.ssh_sessions[ip]
                        if ip in self.ssh_cache:
                            del self.ssh_cache[ip]

                        # 更新UI状态
                        self.update_server_status(item, "未连接", 5)
                        self.append_to_terminal(ip, "\n*** 连接已关闭 ***\n")

                        # 关闭对应的终端标签页
                        for i in range(self.main_window.terminal_execution.terminal_notebook.index("end")):
                            if self.main_window.terminal_execution.terminal_notebook.tab(i, "text") == f"终端-{ip}":
                                self.main_window.terminal_execution.terminal_notebook.forget(i)
                                break

                        self.main_window.show_message(f"已成功关闭 {ip} 服务器连接")
                        logger.info(f"成功关闭 {ip} 服务器连接")

                    except Exception as e:
                        logger.error(f"断开 {ip} 连接时出错: {e}")
                        self.main_window.show_message(f"断开 {ip} 连接时出错: {str(e)}")

    def close_all_connections(self):
        """关闭所有连接（包括缓存）"""
        logger.info("尝试关闭所有连接并清除缓存...")
        # 关闭所有SSH会话
        for ip, session in self.ssh_cache.items():
            try:
                # 1. Close the main SSH connection
                if session['ssh'] is not None:
                    session['ssh'].close()
                    logger.info(f"已关闭 {ip} 的SSH连接")

                # 2. Then, close the jump host connection if it exists
                if 'jump_ssh' in session and session['jump_ssh']:
                    session['jump_ssh'].close()
                    logger.info(f"已关闭 {ip} 的跳板机连接 (在 close_all_connections)")
            except Exception as e:
                logger.warning(f"关闭 {ip} 的SSH连接时出错: {e}")

        # 清空缓存和会话字典
        self.ssh_sessions.clear()
        self.ssh_cache.clear()
        logger.info("已清空SSH会话和缓存")

        # 更新所有服务器状态
        for item in self.server_table.get_children():
            values = list(self.server_table.item(item, "values"))
            if values[5] in ["已连接", "已连接(develop)", "已连接(缓存)", "未连接(已缓存)"]:
                values[5] = "未连接"
                self.server_table.item(item, values=values)

        # 关闭除系统消息外的所有标签页
        for i in range(self.main_window.terminal_execution.terminal_notebook.index("end") - 1, 0, -1):
            self.main_window.terminal_execution.terminal_notebook.forget(i)

        # 清空终端输出字典
        self.terminal_outputs = {}

        self.main_window.show_message("已关闭所有连接并清除缓存")
        logger.info("所有连接已关闭，缓存已清除") 